import org.apache.spark.ml.classification.RandomForestClassifier
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
object breast_cancer {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder().master("local").getOrCreate()
    import spark.implicits._

    val source = spark.read.option("header", "true").option("inferSchema", "true").csv("file:///D:\\IdeaProjects\\machine_learning\\test\\data\\breast_cancer.csv")
//    val source = spark.read.option("header", "true").option("inferSchema", "true").csv("/data/breast_cancer.csv")

    val Array(train,test) = source.randomSplit(Array(0.8, 0.2) )

    val assembler = new VectorAssembler().setInputCols(source.drop("label").columns).setOutputCol("features")

    val rfc = new RandomForestClassifier()

    val model = rfc.fit(assembler.transform(train))

    val predict_train = model.transform(assembler.transform(train))
    val predict = model.transform(assembler.transform(test))

    val accuracy_train = new MulticlassClassificationEvaluator().setMetricName("accuracy").evaluate(predict_train)
    val accuracy_test = new MulticlassClassificationEvaluator().setMetricName("accuracy").evaluate(predict)

    println(s"训练集精确度: $accuracy_train  测试集精确度: $accuracy_test")
  }
}
